import xlrd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
import time

file_path = '22个省市_平均价_猪肉2.xlsx'

work_book = xlrd.open_workbook(file_path)
content = work_book.sheet_by_name('22个省市_平均价_猪肉')

time_list = content.col_values(0)[1:]
pig_price_list = content.col_values(1)[1:]
profit_list = content.col_values(2)[1:]

output_price_list = []
output_profit_list =[]

new_price = 0
new_profit = 0

for pig_price, profit in zip(pig_price_list, profit_list):
    if pig_price != '':
        new_price = pig_price
        output_price_list.append(pig_price)
    else:
        output_price_list.append(new_price)
    if profit != '':
        new_profit = profit
        output_profit_list.append(profit)
    else:
        output_profit_list.append(new_profit)

input_price_list = output_price_list[104:]
input_profit_list = output_profit_list[:-104]

print(len(input_profit_list), len(input_profit_list))

# 绘制散点图
# plt.scatter(input_price_list[350:], input_profit_list[350:])
# plt.show()


regr = LinearRegression()

X = input_profit_list[:350]
Y = input_price_list[:350]

x = np.array(X)
y = np.array(Y)
X = np.c_[np.ones(len(x)), x]
Y = np.c_[y]

#对线性回归实例化
lr = LinearRegression()

#把X,Y放入模型进行训练
lr.fit(X, Y)

#对X数据进行预测
pre_y = lr.predict(X)

#数据可视化
# plt.figure('预测')
# plt.scatter(X[:, -1], Y, c='blue')
# plt.plot(X[:, -1], pre_y, 'r')
# plt.show()
x = [i for i in range(len(Y))]
plt.figure('预测')
plt.plot(x, Y, c='blue')
plt.plot(x, pre_y, 'r')
plt.show()

print('test')
